﻿#include "VtkExportModule.h"
#include "ClassFactory.h"
#include "TimeStep.h"
#include "Analysis.h"
#include "Node.h"
#include "Element.h"
#include "NumericalMethod.h"
#include "mathfem.h"
#include <unordered_set>
REGISTER_ExportModule(VtkExportModule)

VtkExportModule::VtkExportModule(int n, Analysis* e): ExportModule(n, e)
{
	type = NodalRecoveryModel::NodalAveraging;
	smoother = NULL;
}

VtkExportModule::~VtkExportModule()
{
	if (smoother)
	{
		delete smoother;
	}
}


void VtkExportModule::initial(Tag& tag)
{
	std::list<std::string> dofStrings;
	++tag;
	tag.getValue(dofStrings);
	dofVarsToExport.reserve(20);
	intVarsToExport.reserve(20);
	cellVarsToExport.reserve(20);
	for (auto& istr : dofStrings)
	{
		if (strcmp(istr.c_str(), "displacement") == 0)
		{
			dofVarsToExport.push_back(Displacement);
		}
		else if (strcmp(istr.c_str(), "stress") == 0)
		{
			intVarsToExport.push_back(StressIVS);
		}
		else if (strcmp(istr.c_str(), "strain") == 0)
		{
			intVarsToExport.push_back(StrainIVS);
		}
		else if (strcmp(istr.c_str(), "estress") == 0)
		{
			cellVarsToExport.push_back(StressIVS);
		}
		else if (strcmp(istr.c_str(), "estrain") == 0)
		{
			cellVarsToExport.push_back(StrainIVS);
		}
		else if (strcmp(istr.c_str(), "peeq") == 0)
		{
			intVarsToExport.push_back(EPSIVS);
		}
		else if (strcmp(istr.c_str(), "epeeq") == 0)
		{
			cellVarsToExport.push_back(EPSIVS);
		}
	}
}

void VtkExportModule::doOutput(TimeStep* ctime, bool force)
{
	if (!(TestTimeStepOut(ctime) || force))
	{
		return;
	}
	FILE* file = this->getOutputStream(ctime);
	// 输出vtk前面的注释
	fprintf(file, "# vtk DataFile Version 2.0\n");
	fprintf(file, "Output for timestep %d time %f\n", ctime->GetNumber(), ctime->GetTargetTime());
	fprintf(file, "ASCII\n");
	fprintf(file, "DATASET UNSTRUCTURED_GRID\n");

	// 输出节点坐标(是全部输出的)
	Domain* d = analysis->GetDomain(1);
	DoubleArray* coords;
	int i, inode, nnodes = d->GetNumberOfNode();
	this->getSmoother();
	fprintf(file, "POINTS %d double\n", nnodes);
	for (inode = 1; inode <= nnodes; inode++)
	{
		coords = d->GetNode(inode)->GetCoords();
		for (i = 1; i <= coords->getSize(); i++)
		{
			fprintf(file, "%e ", coords->at(i - 1));
		}
		for (i = coords->getSize() + 1; i <= 3; i++)
		{
			fprintf(file, "%e ", 0.0);
		}
		fprintf(file, "\n");
	}

	// 输出单元节点集
	int count = 0;
	int cellsize = 0, ncell;
	for (auto& elem : d->GetElements())
	{
		// 显示所有单元，只要有网格就进行显示
		count++;
		ncell = this->getNumberOfElementCells(elem.get());
		cellsize += ncell + ncell * this->getNumerOfNodesPerCell(this->getVtkCellType(elem.get()));
	}
	for (auto& ir : d->getNumericalMethodList())
	{
		if (ir->getItsOutDataType() != NMOD_POINT || !ir->isActived())
		{
			// 基于离散节点生成的积分点类方法才会进行显示
			continue;
		}
		for (auto& ip : *ir)
		{
			count++;
			cellsize += 2;
		}
	}
	fprintf(file, "CELLS %d %d\n", count, cellsize);
	int vtkcelltype, node;
	IntArray cellNodes;
	for (auto& elem : d->GetElements())
	{
		vtkcelltype = this->getVtkCellType(elem.get());
		node = this->getNumerOfNodesPerCell(vtkcelltype);
		this->getVtkCellNodes(cellNodes, elem.get());
		fprintf(file, "%d ", node);
		for (i = 1; i <= node; i++)
		{
			fprintf(file, "%d ", cellNodes(i) - 1);
		}
		fprintf(file, "\n");
	}
	// 单元类型
	fprintf(file, "CELL_TYPES %d\n", count);
	for (auto& elem : d->GetElements())
	{
		vtkcelltype = this->getVtkCellType(elem.get());
		fprintf(file, "%d\n", vtkcelltype);
	}
	for (auto& ir : d->getNumericalMethodList())
	{
		if (ir->getItsOutDataType() != NMOD_POINT || !ir->isActived())
		{
			// 非粒子法的数值方法跳过，只有粒子法的点才会显示云图
			continue;
		}
		for (auto& ip : *ir)
		{
			fprintf(file, "1\n"); // point 类型
		}
	}
	// 数据输出
	// 单元数据输出
	if (cellVarsToExport.getSize())
	{
		this->exportCellVars(file, count, ctime);
	}

	if (dofVarsToExport.getSize() || intVarsToExport.getSize())
	{
		// 默认输出所有节点
		fprintf(file, "POINT_DATA %d\n", nnodes);
	}
	// 输出节点自由度上的值，这里没哟积分点上的值，不需要进行积分到节点的变量转换
	this->exportDofVars(file, ctime);
	// 输出节点上的值（积分点到节点上的转换）
	this->exportIntVars(file, ctime);

	fclose(file);
}


void VtkExportModule::exportIntVars(FILE* stream, TimeStep* ctime)
{
	/*
	* 暂且不考虑积分点输出位移
	*/
	int size = intVarsToExport.getSize();
	if (size == 0)
	{
		return;
	}

	InternalType it;
	// 单元集合准备
	Domain* d = analysis->GetDomain(1);
	int nnode = d->GetNumberOfNode();

	// 创建集合
	IntArray elemSet;
	elemSet.assignfrom(d->GetNumberOfElement());


	size = intVarsToExport.getSize();
	IntArray regionSizes(size);
	for (int i = 1; i <= size; i++)
	{
		it = GetInternalType((InternalValueStateType)intVarsToExport(i));
		regionSizes(i) = getSizeOfInternalType(it);
	}

	// 创建分析的
	FEOutManager* OutOthers = new FEOutManager;
	OutOthers->createOutModules(intVarsToExport, nnode);

	// 磨平器
	this->getSmoother();
	smoother->recoverNodalValues(elemSet, intVarsToExport, regionSizes, ctime);

	int countsize = 0;
	for (int i = 1; i <= size; i++)
	{
		InternalValueStateType valtype = (InternalValueStateType)intVarsToExport(i);
		it = GetInternalType(valtype);
		if (it == SCALARIT)
		{
			fprintf(stream, "SCALARS %s double 1\n", InternalValueTypeToString(valtype).c_str());
		}
		else if (it == VECTORIT)
		{
			fprintf(stream, "VECTORS %s double\n", InternalValueTypeToString(valtype).c_str());
		}
		else if (it == TENSORIT || it == TENSOREIT || it == TENSOREGIT)
		{
			fprintf(stream, "TENSORS %s double\n", InternalValueTypeToString(valtype).c_str());
		}
		else
		{
			ERROR("unspport type %s \n", InternalValueTypeToString(valtype).c_str());
		}
		if (it == SCALARIT)
		{
			fprintf(stream, "LOOKUP_TABLE default\n");
		}
		const DoubleArray* val;
		DoubleArray ival;
		mat3d t;

		if (i > 1)
		{
			countsize += regionSizes(i - 1);
		}
		for (int inode = 1; inode <= nnode; inode++)
		{
			this->smoother->getNodalVector(val, inode);
			if (val == NULL)
			{
				int jsize = 0;
				for (int m = 1; m <= regionSizes.getSize(); ++m)
				{
					jsize += regionSizes(i);
				}
				ival.resize(jsize);
				ival.zero();
				val = &ival;
			}
			if (it == SCALARIT)
			{
				fprintf(stream, "%e ", val->at(countsize));
			}
			else if (it == VECTORIT)
			{
				// 忽略节点坐标系
				fprintf(stream, "%e %e %e ", val->at(countsize), val->at(countsize + 1), val->at(countsize + 2));
			}
			else if (it == TENSORIT || it == TENSOREIT)
			{
				// 对称张量 这里感觉是不正确的如果是TENSOREIT， 3,4,5是需要除以2的(这里有着应变张量转换问题)


				t(0, 0) = val->at(countsize); //xx
				t(1, 1) = val->at(countsize + 1); //yy
				t(2, 2) = val->at(countsize + 2); //zz
				t(1, 2) = val->at(countsize + 3); //yz
				t(2, 1) = val->at(countsize + 3); //zy
				t(0, 2) = val->at(countsize + 4); //xz
				t(2, 0) = val->at(countsize + 4); //zx
				t(0, 1) = val->at(countsize + 5); //xy
				t(1, 0) = val->at(countsize + 5); //yx

				for (int ii = 0; ii < 3; ii++)
				{
					for (int jj = 0; jj < 3; jj++)
					{
						fprintf(stream, "%e ", t(ii, jj));
					}
				}
				OutOthers->setVariables(t, valtype, inode);
			}
			else if (it == TENSOREGIT)
			{
				t(0, 0) = val->at(countsize);
				t(1, 1) = val->at(countsize + 1);
				t(2, 2) = val->at(countsize + 2);
				t(1, 2) = val->at(countsize + 3);
				t(2, 1) = val->at(countsize + 4);
				t(0, 2) = val->at(countsize + 5);
				t(2, 0) = val->at(countsize + 6);
				t(0, 1) = val->at(countsize + 7);
				t(1, 0) = val->at(countsize + 8);
				for (int ii = 0; ii < 3; ii++)
				{
					for (int jj = 0; jj < 3; jj++)
					{
						fprintf(stream, "%e ", t(ii, jj));
					}
				}
				OutOthers->setVariables(t, valtype, inode);
			}
			fprintf(stream, "\n");
		}
	}

	// 对于部分张量需要输出该张量其他东西（比如应力需要其等效应力、剪应力等）
	for (int i = 0; i < size; ++i)
	{
		InternalValueStateType valtype = (InternalValueStateType)intVarsToExport[i];

		// 获取该张量的相关其他输出
		FEOut* iout = OutOthers->getFEOut(valtype);
		if (!iout)
		{
			continue;
		}

		// 获取该张量其他尺寸大小
		int nsize = iout->MagSize();
		for (int j = 1; j <= nsize; ++j)
		{
			const DoubleArray& val = iout->getMag(j);
			fprintf(stream, "SCALARS %s double 1\n", iout->getMagString(j));
			fprintf(stream, "LOOKUP_TABLE default\n");

			// @TODO节点循环(先这样，可能有更好的做法)
			for (int inode = 0; inode < nnode; inode++)
			{
				fprintf(stream, "%e\n ", val[inode]);
			}
			fprintf(stream, "\n");
		}
	}
	delete OutOthers;
}

void VtkExportModule::exportCellVars(FILE* stream, int elemToProcess, TimeStep* ctime)
{
	Domain* d = analysis->GetDomain(1);
	InternalValueStateType type;
	DoubleArray help, vec;
	mat3d t;
	int jsize;
	double gptot = 0.;
	int nelem = d->GetNumberOfElement();
	mat3d s;
	// 创建分析的
	FEOutManager* OutOthers = new FEOutManager;
	OutOthers->createOutModules(cellVarsToExport, nelem);
	fprintf(stream, "\nCELL_DATA %d\n", elemToProcess);

	for (int i = 1; i <= cellVarsToExport.getSize(); i++)
	{
		type = (InternalValueStateType)cellVarsToExport(i);
		InternalType it = GetInternalType(type);
		FEOut* iout = OutOthers->getFEOut(type);

		if (it == SCALARIT)
		{
			fprintf(stream, "SCALARS %s double\n LOOKUP_TABLE default\n", InternalValueTypeToString(type).c_str());
		}
		else if (it == TENSORIT || it == TENSOREIT)
		{
			fprintf(stream, "TENSORS %s double\n", InternalValueTypeToString(type).c_str());
		}
		else
		{
			fprintf(stream, "VECTORS %s double\n", InternalValueTypeToString(type).c_str());
		}
		int count = 1;
		for (auto& elem : d->GetElements())
		{
			if (elem->isNumericalMethodOn())
			{
				continue;
			}

			gptot = 0.;
			vec.clear();
			double* w = elem->GaussWeights();
			for (int n = 0; n < elem->numberOfGaussPoints; ++n)
			{
				IntegrationPoint* ip = elem->getFemRule()->getIntegrationPoint(n + 1);
				elem->getIPValue(help, ip, type, ctime);
				gptot += w[n];
				vec.Add(w[n], help);
			}
			vec.Scale(1 / gptot);


			if (it == SCALARIT)
			{
				if (vec.getSize())
				{
					fprintf(stream, "%e ", vec(1));
				}
				else
				{
					fprintf(stream, "%e ", 0.);
				}
			}
			else if (it == VECTORIT)
			{
				int size = min_fem(3, vec.getSize());
				for (jsize = 1; jsize <= size; jsize++)
				{
					fprintf(stream, "%e ", vec(jsize));
				}
				for (jsize = size + 1; jsize <= 3; jsize++)
				{
					fprintf(stream, "%e ", 0.);
				}
			}
			else if (it == TENSORIT)
			{
				t(0, 0) = vec(1);
				t(1, 1) = vec(2);
				t(2, 2) = vec(3);
				t(1, 2) = vec(4);
				t(2, 1) = vec(4);
				t(0, 2) = vec(5);
				t(2, 0) = vec(5);
				t(0, 1) = vec(6);
				t(1, 0) = vec(6);
				for (int ii = 0; ii < 3; ii++)
				{
					for (int jj = 0; jj < 3; jj++)
					{
						fprintf(stream, "%e ", t(ii, jj));
					}
				}
				if (iout)
				{
					iout->setVals(t, count);
				}
			}
			else
			{
				t(0, 0) = vec(1);
				t(1, 1) = vec(2);
				t(2, 2) = vec(3);
				t(1, 2) = vec(4);
				t(2, 1) = vec(4);
				t(0, 2) = vec(5);
				t(2, 0) = vec(5);
				t(0, 1) = vec(6);
				t(1, 0) = vec(6);
				for (int ii = 1; ii <= 3; ii++)
				{
					for (int jj = 1; jj <= 3; jj++)
					{
						fprintf(stream, "%e ", t(ii, jj));
					}
				}
				iout->setVals(t, count);
			}
			++count;
			fprintf(stream, "\n");
		}
		// 对于部分张量需要输出该张量其他东西（比如应力需要其等效应力、剪应力等）
		for (int j = 0; j < cellVarsToExport.getSize(); ++j)
		{
			InternalValueStateType valtype = (InternalValueStateType)cellVarsToExport[j];

			// 获取该张量的相关其他输出
			FEOut* iout = OutOthers->getFEOut(valtype);
			if (!iout)
			{
				continue;
			}

			// 获取该张量其他尺寸大小
			int nsize = iout->MagSize();
			for (int k = 1; k <= nsize; ++k)
			{
				const DoubleArray& val = iout->getMag(k);
				fprintf(stream, "SCALARS ELEM_%s double 1\n", iout->getMagString(k));
				fprintf(stream, "LOOKUP_TABLE default\n");

				// @TODO单元循环(先这样，可能有更好的做法)
				for (int ielem = 0; ielem < val.getSize(); ielem++)
				{
					fprintf(stream, "%e\n ", val[ielem]);
				}
				fprintf(stream, "\n");
			}
		}
	}
	//打印到文件
	//OutOthers->print2File("Element");
	delete OutOthers;
}


void VtkExportModule::initialize()
{
	ExportModule::initialize();
	if (this->smoother)
	{
		delete this->smoother;
		smoother = NULL;
	}
	//OutOthers = new FEOutManager;
}

int VtkExportModule::getNumerOfNodesPerCell(int celltype)
{
	switch (celltype)
	{
	case 1:
		return 1;
	case 3:
		return 2;

	case 5:
	case 21:
		return 3;

	case 9:
	case 10:
		return 4;

	case 14:
		return 5;

	case 13:
	case 22:
		return 6;

	case 12:
	case 23:
		return 8;

	case 24:
		return 10;

	case 25:
		return 20;

	case 26:
		return 15;
	default:
		FATALERROR("vtk不支持的单元类型");
	}
}

FILE* VtkExportModule::getOutputStream(TimeStep* ctime)
{
	std::string name;
	char txt[30];
	sprintf(txt, "_%d", ctime->GetNumber());
	name = this->GetOutputBaseFileName() + txt + ".vtk";
	FILE* answer = fopen(name.c_str(), "w");
	if (answer == NULL)
	{
		FATALERROR("打开文件%s失败！", name.c_str());
	}
	return answer;
}

int VtkExportModule::getNumberOfElementCells(Element* cell)
{
	ElementGeoType elemGT = cell->getElementGeoType();
	if ((elemGT == POINT_EGT) || (elemGT == LINE_1_EGT) ||
		(elemGT == LINE_2_EGT) || (elemGT == TRI_1_EGT) ||
		(elemGT == TRI_2_EGT) || (elemGT == TETRA_1_EGT) ||
		(elemGT == TETRA_2_EGT) || (elemGT == QUAD_1_EGT) ||
		(elemGT == QUAD_2_EGT) || (elemGT == HEX_1_EGT) ||
		(elemGT == HEX_2_EGT) || (elemGT == WEDGE_1_EGT) ||
		(elemGT == WEDGE_2_EGT) || elemGT == PYRAM_1_EGT)
	{
		return 1;
	}
	else
	{
		return 0;
		//ERROR("unsupported element geometry type");
	}
}

int VtkExportModule::getVtkCellType(Element* cell)
{
	ElementGeoType elemGT = cell->getElementGeoType();
	int vtkCellType;
	if (elemGT == POINT_EGT)
	{
		vtkCellType = 1;
	}
	else if (elemGT == LINE_1_EGT)
	{
		vtkCellType = 3;
	}
	else if (elemGT == LINE_2_EGT)
	{
		vtkCellType = 21;
	}
	else if (elemGT == TRI_1_EGT)
	{
		vtkCellType = 5;
	}
	else if (elemGT == TRI_2_EGT)
	{
		vtkCellType = 22;
	}
	else if (elemGT == TETRA_1_EGT)
	{
		vtkCellType = 10;
	}
	else if (elemGT == TETRA_2_EGT)
	{
		vtkCellType = 24;
	}
	else if (elemGT == QUAD_1_EGT)
	{
		vtkCellType = 9;
	}
	else if (elemGT == QUAD_2_EGT)
	{
		vtkCellType = 23;
	}
	else if (elemGT == HEX_1_EGT)
	{
		vtkCellType = 12;
	}
	else if (elemGT == HEX_2_EGT)
	{
		vtkCellType = 25;
	}
	else if (elemGT == WEDGE_1_EGT)
	{
		vtkCellType = 13;
	}
	else if (elemGT == WEDGE_2_EGT)
	{
		vtkCellType = 26;
	}
	else if (elemGT == PYRAM_1_EGT)
	{
		vtkCellType = 14;
	}

#ifdef _DEBUG
	else
	{
		DEBUG_ERROR("不支持的单元几何类型");
	}
#endif

	return vtkCellType;
}

void VtkExportModule::getVtkCellNodes(IntArray& answer, Element* elem)
{
	ElementGeoType elemGT = elem->getElementGeoType();
	int i, nelemNodes;

	if ((elemGT == POINT_EGT) ||
		(elemGT == LINE_1_EGT) || (elemGT == LINE_2_EGT) ||
		(elemGT == TRI_1_EGT) || (elemGT == TRI_2_EGT) ||
		(elemGT == TETRA_1_EGT) || (elemGT == TETRA_2_EGT) ||
		(elemGT == QUAD_1_EGT) || (elemGT == QUAD_2_EGT) ||
		(elemGT == HEX_1_EGT) || (elemGT == HEX_2_EGT) ||
		(elemGT == WEDGE_1_EGT) || (elemGT == WEDGE_2_EGT) || (elemGT == PYRAM_1_EGT))
	{
		nelemNodes = elem->GetNumberOfNodes();
		answer.resize(nelemNodes);
		for (i = 1; i <= nelemNodes; i++)
		{
			answer(i) = elem->GetNode(i)->GetNumber();
		}
	}
#ifdef _DEBUG
	else
	{
		DEBUG_ERROR("尚未支持!");
	}
#endif
}

void VtkExportModule::exportDofVars(FILE* stream, TimeStep* ctime)
{
	for (auto& var : dofVarsToExport)
	{
		DofValueType type = (DofValueType)var;
		this->exportDofVarsAs(type, stream, ctime);
	}
}

void VtkExportModule::exportDofVarsAs(DofValueType type, FILE* stream, TimeStep* ctime)
{
	InternalType it;
	int scalars = 1;
	if (type == Displacement || type == Velocity || type == EigenVector)
	{
		it = VECTORIT;
	}
	else if (type == Temperature)
	{
		it = SCALARIT;
	}
#ifdef _DEBUG
	else
	{
		DEBUG_ERROR("尚未支持 %s yet!", DofValueTypeToString(type));
	}
#endif

	if (it == SCALARIT)
	{
		fprintf(stream, "SCALARS %s double %d\n", DofValueTypeToString(type), scalars);
	}
	else if (it == VECTORIT)
	{
		fprintf(stream, "VECTORS %s double\n", DofValueTypeToString(type));
	}
	else
	{
		ERROR("unspport variable type\n");
	}

	if (it == SCALARIT)
	{
		fprintf(stream, "LOOKUP_TABLE default\n");
	}

	// 所有节点进行输出，因为这是自由度上的不是积分点上的变量
	Domain* d = analysis->GetDomain(1);
	DoubleArray val;
	DofType id;
	int size, j;
	for (auto& inode : d->GetNodes())
	{
		if (type == Displacement || type == EigenVector || type == Velocity)
		{
			val.resize(3);
			val.zero();
			for (Dof* dof : *inode)
			{
				id = dof->getDofType();
				if (id == DispX || id == VelX)
				{
					val(1) = dof->getUpdateValue(FieldTotalVal, ctime);
				}
				else if (id == DispY || id == VelY)
				{
					val(2) = dof->getUpdateValue(FieldTotalVal, ctime);
				}
				else if (id == DispZ || id == VelZ)
				{
					val(3) = dof->getUpdateValue(FieldTotalVal, ctime);
				}
			}
		}
		else if (type == Temperature)
		{
			val.resize(1);
			val.zero();
			for (Dof* dof : *inode)
			{
				id = dof->getDofType();
				if (id == TField)
				{
					val(1) = dof->getUpdateValue(FieldTotalVal, ctime);
				}
			}
		}
		else
		{
			ERROR("unsupport DofvalueType %s", DofValueTypeToString(type));
		}
		if (it == SCALARIT)
		{
			if (val.getSize())
			{
				for (j = 1; j <= scalars; j++)
				{
					fprintf(stream, "%e ", val(j));
				}
			}
			else
			{
				fprintf(stream, "%e ", 0.0);
			}
		}
		else if (it == VECTORIT)
		{
			size = min_fem(3, val.getSize());
			// 忽略节点坐标系
			for (j = 1; j <= size; j++)
			{
				fprintf(stream, "%e ", val(j));
			}
			for (j = size + 1; j <= 3; j++)
			{
				fprintf(stream, "%e ", 0.0);
			}
		}
		fprintf(stream, "\n");
	}
	fprintf(stream, "\n");
}

NodalRecoveryModel* VtkExportModule::getSmoother()
{
	if (smoother == NULL)
	{
		smoother = classFactory.CreateNodalRecoveryModel(type, analysis->GetDomain(1));
	}
	return this->smoother;
}
